
# 导入数学库，提供 sqrt、pow 等基础数学函数
import math

# 导入随机库，支持随机抽样与随机数生成
import random
from collections import defaultdict

# 计算两个向量之间的欧氏距离(L2距离)
def l2_distance(vec1, vec2):
    return math.sqrt(sum(pow(a - b, 2) for a, b in zip(vec1, vec2)))

def kmeans(vectors, nlist):
    # 随机选择初始中心点
    centers = [random.choice(vectors)]
    print("Initial center:", centers)

    # 解释：下面这段while是K-Means++的“轮盘赌”抽样逻辑
    # 步骤：计算每个点到最近中心的距离平方 -> 得到概率权重 -> 依概率抽样新的中心点

    # 当中心数量不足k个时，继续选择新的中心点
    while len(centers) < nlist:
        # 第一步：每个数据点距离最近中心的平方，作为该点的权重
        dists = [min(l2_distance(vec, center) ** 2 for center in centers) for vec in vectors]
        # for vector in vectors:
        #     print("Distances to centers:", min([l2_distance(vector, center) ** 2 for center in centers]))
        # print("Distances to nearest center:", dists)
        # 第二步：对所有权重求和，得到轮盘赌的总长度
        total = sum(dists)
        # 第三步：在[0, total]区间上随机抽一个阈值r，当作轮盘指针
        r = random.random() * total

        print('dists', dists)
        print("Total distance sum:", total)
        print("Random threshold r:", r)
        print("centers:", centers)

        # prefix 累加当前遍历到的权重，用于定位r落在哪个区间
        prefix = 0.0

        # 遍历所有点和对应权重，寻找累积首次超过r的那个点
        for v, d in zip(vectors, dists):
            prefix += d
            # 一旦累积权重超过r，说明当前点被选中作为新中心
            if prefix >= r:
                centers.append(v)
                break

    print("Final centers:", centers)
    # 初始化一个“簇字典”，key是中心索引，value是该中心点下的样本列表
    clusters = {i: [] for i in range(nlist)}
    print("Clusters initialized:", clusters)
    # 遍历所有向量，找到距离最近的中心，将该向量归入对应簇
    for v in vectors:
        idx = min(range(nlist), key=lambda i: l2_distance(v, centers[i]))
        # print(f"Vector {v} assigned to cluster {idx}")
        clusters[idx].append(v)
    print("遍历之后的簇字典:", clusters)

    
    # 对每个非空簇计算均值，作为中心的最新位置
    for i, points in clusters.items():
        if points:
            # zip(*points) 将点按维度拆开，sum(dim)/len(points) 是维度均值
            centers[i] = [sum(dim) / len(points) for dim in zip(*points)]
    print("Updated centers:", centers)
    return centers

def build_ivf_flat(vectors, nlist):
    # 先通过K-Means算法找到nlist个中心点
    centers = kmeans(vectors, nlist)
    # defaultdict(list) 便于向桶里 append，无需检查 key 是否存在
    buckets = defaultdict(list)

    # 遍历所有向量，依据最近中心编号将其放入对应桶
    for vec in vectors:
        idx = min(range(nlist), key=lambda i: l2_distance(vec, centers[i]))
        buckets[idx].append(vec)

    # 返回构建好的中心列表与倒排结构
    return centers, buckets

# 在 IVF-Flat 索引上执行查询，近似返回 topk 个最近向量
def search_ivf_flat(query, centers, buckets, nprobe, topk=3):
    # 先按查询向量到各中心的距离排序，得到“探查顺序”
    probe_order = sorted(range(len(centers)), key=lambda i: l2_distance(query, centers[i]))
    
    # 用列表保存候选 (距离, 向量) 对
    candidate_pairs = []

    # 只深入距离最近的 nprobe 个中心对应的桶
    for idx in probe_order[:nprobe]:
        for vec in buckets[idx]:
            candidate_pairs.append((l2_distance(query, vec), vec))
    
    # 候选列表按距离升序排列，便于截取前 topk 个
    candidate_pairs.sort(key=lambda x: x[0])

    # 返回最接近查询的 topk 个候选
    return candidate_pairs[:topk]

# 主程序：演示 IVF-Flat 的构建与查询流程
def main():
    # 设定随机种子，保证示例结果可复现
    random.seed(0)

    # 构造 8 个二维向量作为“数据库”
    database = [
        [0.1, 0.2],
        [0.15, 0.18],
        [0.8, 0.85],
        [0.82, 0.79],
        [0.4, 0.4],
        [0.42, 0.45],
        [0.9, 0.1],
        [0.88, 0.15],
    ]

    # 查询向量：我们要找与它相似的数据库向量
    query = [0.12, 0.22]

    # nlist 表示建索引时划分的中心（簇）个数
    nlist = 4

    # nprobe 表示查询时会探查多少个最近的中心
    nprobe = 2

    # 构建 IVF-Flat 索引，得到中心与倒排桶
    centers, buckets = build_ivf_flat(database, nlist=nlist)

    # 执行检索，返回距离最近的 top3 个结果
    results = search_ivf_flat(query, centers, buckets, nprobe=nprobe, topk=3)

    # 打印本次检索的参数设置
    print(f"nlist={nlist}, nprobe={nprobe}")

    # 打印检索到的候选向量及其距离
    print("最近向量：")
    for dist, vec in results:
        print(f"向量={vec}, 距离={dist:.4f}")

if __name__ == "__main__":
    main()